import os, sys
import os.path as osp
import numpy as np
import cv2
import json

from tool import filesystem
from paddlex import create_model

seal_cls_model = create_model(model_name="PP-LCNet_x1_0_pedestrian_attribute", model_dir="./model/seal_cls")



def seal_class_func(image, debug=0):
    output = seal_cls_model.predict(input=image, batch_size=1)
    outputs = list(output)
    if debug:
        for res in outputs:
            # res.print()
            res.save_to_img(save_path="./output/")
            res.save_to_json(save_path="./output/")

    ret_data = {}
    ret_data["label_names"] = []
    ret_data["scores"] = []
    ret_data["class_ids"] = []
    if len(outputs) == 0:
        return ret_data
    
    json_rec_data = outputs[0].json["res"]
    ret_data["label_names"] = json_rec_data["label_names"]
    ret_data["scores"] = json_rec_data["scores"]
    ret_data["class_ids"] = json_rec_data["class_ids"]
    return ret_data


if __name__ == "__main__":

    file_path = sys.argv[1]
    file_paths = []
    if osp.isfile(file_path):
        file_paths.append(file_path)
    else:
        file_paths = filesystem.get_all_filepath(file_path, [".jpg", ".png"])
    
    for file_path in file_paths:
        bgr_img = cv2.imread(file_path)
            
        res = seal_class_func(bgr_img, debug=1)
        # res = handwrite_rec_func(bgr_img, debug=1)
        print(osp.basename(file_path), res)

"""
完整的类别
    labels = [
        "印章颜色红色",
        "印章颜色蓝色",
        "印章颜色黑色",
        "缺失程度严重",
        "缺失程度轻微",
        "模糊程度严重",
        "模糊程度轻微",
        "公章",
        # "财务章",
        # "合同章",
        # "发票章",
        # "法人代表人名章",
        "重叠",
        "中心标志五角星",
        "中心标志党徽"
    ]
"""
